use actix_web::{get, post, web, HttpRequest, Responder};
use actix_ws::Message;
use futures_util::TryStreamExt;
use reqwest::Client;
use serde_json::Value;

#[get("/ws")]
pub async fn ws(req: HttpRequest, body: web::Payload) -> actix_web::Result<impl Responder> {
    let (response, mut session, mut msg_stream) = actix_ws::handle(&req, body)?;
    let client = Client::new();

    actix_web::rt::spawn(async move {
        while let Some(Ok(msg)) = msg_stream.recv().await {
            match msg {
                Message::Ping(bytes) => {
                    if session.pong(&bytes).await.is_err() {
                        return;
                    }
                }
                Message::Text(msg) => {
                    println!("Got text: {}", msg);
                    let mut stream = client
                        .post("http://127.0.0.1:11434/api/generate")
                        .json(&serde_json::json!({ "model": "llama3.1", "prompt": *msg, "stream": true}))
                        .send()
                        .await
                        .expect("failed to send")
                        .bytes_stream();

                    while let Ok(item) = stream.try_next().await {
                        if let Some(chunk) = item {
                            let chunk_str =
                                String::from_utf8(chunk.to_vec()).expect("Invalid UTF-8");

                            let json_chunk: Value =
                                serde_json::from_str(&chunk_str).expect("parse json error");

                            if let Some(response) =
                                json_chunk.get("response").and_then(|v| v.as_str())
                            {
                                println!("Response: {}", response);
                                // Send the chunk back to the client
                                if session.text(response).await.is_err() {
                                    break;
                                }
                            }
                        }
                    }
                }
                _ => break,
            }
        }

        let _ = session.close(None).await;
    });

    Ok(response)
}
